import logging
from typing import List, Optional, Dict
from collections import deque
from datetime import datetime, timezone

from .base_memory import BaseMemory
from ..agent_types import MemoryRecord

logger = logging.getLogger(__name__)

class ShortTermMemory(BaseMemory):
    """
    Simple in-memory implementation of short-term (working) memory using a deque.
    Stores recent records for the current task context.
    """

    def __init__(self, max_records: int = 100):
        self.max_records = max_records
        # Store records per task ID to keep contexts separate (optional)
        # If not per-task, just use one deque.
        self.task_memory: Dict[str, deque[MemoryRecord]] = {}
        self.default_memory: deque[MemoryRecord] = deque(maxlen=max_records) # For records without task_id
        logger.info(f"Initialized ShortTermMemory with max_records={max_records}")

    def _get_memory_deque(self, task_id: Optional[str]) -> deque[MemoryRecord]:
         """Gets the appropriate deque based on task_id."""
         if task_id:
             if task_id not in self.task_memory:
                 self.task_memory[task_id] = deque(maxlen=self.max_records)
             return self.task_memory[task_id]
         else:
             return self.default_memory

    async def add_record(self, record: MemoryRecord, task_id: Optional[str] = None) -> None:
        """Adds a record to the appropriate deque."""
        # Ensure timestamp is set
        if not record.timestamp: record.timestamp = datetime.now(timezone.utc)
        memory = self._get_memory_deque(task_id or record.related_task_id)
        memory.append(record)
        logger.debug(f"Added record type '{record.type}' to short-term memory (Task: {task_id or 'Default'}).")

    async def add_records(self, records: List[MemoryRecord], task_id: Optional[str] = None) -> None:
        """Adds multiple records."""
        memory = self._get_memory_deque(task_id or (records[0].related_task_id if records else None))
        for record in records:
             if not record.timestamp: record.timestamp = datetime.now(timezone.utc)
             memory.append(record)
        logger.debug(f"Added {len(records)} records to short-term memory (Task: {task_id or 'Default'}).")


    async def retrieve_relevant_records(
        self,
        query: str, # Simple implementation might ignore query and just return recent
        task_id: Optional[str] = None,
        limit: int = 10,
        filter_criteria: Optional[dict] = None # Basic filtering by type supported
    ) -> List[MemoryRecord]:
        """Retrieves recent records, optionally filtered by type. Ignores query text."""
        # TODO: Implement more sophisticated relevance filtering if needed.
        memory = self._get_memory_deque(task_id)
        record_type = filter_criteria.get("type") if filter_criteria else None

        # Iterate from most recent (right side of deque)
        relevant_records = []
        count = 0
        for record in reversed(memory):
             if count >= limit: break
             if record_type and record.type != record_type: continue
             # Add more complex filtering here if needed
             relevant_records.append(record)
             count += 1

        logger.debug(f"Retrieved {len(relevant_records)} relevant records (Type: {record_type or 'Any'}, Task: {task_id or 'Default'}). Query ignored.")
        return list(reversed(relevant_records)) # Return in chronological order

    async def get_recent_records(
        self,
        task_id: Optional[str] = None,
        limit: int = 10,
        record_type: Optional[str] = None
    ) -> List[MemoryRecord]:
         """Gets the most recent records, optionally filtered by type."""
         memory = self._get_memory_deque(task_id)
         # Iterate from most recent (right side of deque)
         recent_records = []
         count = 0
         for record in reversed(memory):
             if count >= limit: break
             if record_type and record.type != record_type: continue
             recent_records.append(record)
             count += 1
         logger.debug(f"Retrieved {len(recent_records)} recent records (Type: {record_type or 'Any'}, Task: {task_id or 'Default'}).")
         return list(reversed(recent_records)) # Return in chronological order

    async def clear_task_memory(self, task_id: str) -> None:
        """Clears memory for a specific task."""
        if task_id in self.task_memory:
            self.task_memory[task_id].clear()
            del self.task_memory[task_id] # Remove deque itself
            logger.info(f"Cleared short-term memory for Task ID: {task_id}")

    async def clear_all_memory(self) -> None:
        """Clears all short-term memory."""
        self.task_memory.clear()
        self.default_memory.clear()
        logger.info("Cleared all short-term memory.")